#!/usr/bin/env python
# Copyright 1999-2024 Gentoo Authors
# Distributed under the terms of the GNU General Public License v2

import argparse
import errno
import math
import signal
import subprocess
import sys

from os import path as osp

if osp.isfile(
    osp.join(osp.dirname(osp.dirname(osp.realpath(__file__))), ".portage_not_installed")
):
    sys.path.insert(
        0, osp.join(osp.dirname(osp.dirname(osp.realpath(__file__))), "lib")
    )
import portage

portage._internal_caller = True
from portage import os
from portage import xpak, gpkg
from portage.const import SUPPORTED_GENTOO_BINPKG_FORMATS
from portage.dbapi.dep_expand import dep_expand
from portage.dep import Atom, use_reduce
from portage.exception import (
    AmbiguousPackageName,
    CompressorOperationFailed,
    InvalidAtom,
    InvalidData,
    InvalidBinaryPackageFormat,
    InvalidDependString,
    PackageSetNotFound,
    PermissionDenied,
)
from portage.util import ensure_dirs, shlex_split, varexpand, _xattr
from portage.util.cpuinfo import makeopts_to_job_count

xattr = _xattr.xattr
from portage._sets import load_default_config, SETPREFIX
from portage.process import find_binary
from portage.util.compression_probe import _compressors
from portage.util._eventloop.global_event_loop import global_event_loop
from portage.gpg import GPG


def quickpkg_atom(options, infos, arg, eout):
    settings = portage.settings
    eroot = portage.settings["EROOT"]
    trees = portage.db[eroot]
    vartree = trees["vartree"]
    vardb = vartree.dbapi
    bintree = trees["bintree"]

    include_config = options.include_config == "y"
    include_unmodified_config = options.include_unmodified_config == "y"
    fix_metadata_keys = ["PF", "CATEGORY"]

    try:
        atom = dep_expand(arg, mydb=vardb, settings=vartree.settings)
    except AmbiguousPackageName as e:
        # Multiple matches thrown from cpv_expand
        eout.eerror(f"Please use a more specific atom: {' '.join(e.args[0])}")
        del e
        infos["missing"].append(arg)
        return 1
    except (InvalidAtom, InvalidData):
        eout.eerror(f"Invalid atom: {arg}")
        infos["missing"].append(arg)
        return 1
    if atom[:1] == "=" and arg[:1] != "=":
        # dep_expand() allows missing '=' but it's really invalid
        eout.eerror(f"Invalid atom: {arg}")
        infos["missing"].append(arg)
        return 1

    matches = vardb.match(atom)
    pkgs_for_arg = 0
    retval = 0
    for cpv in matches:
        excluded_config_files = []
        dblnk = vardb._dblink(cpv)
        have_lock = False

        if "__PORTAGE_INHERIT_VARDB_LOCK" not in settings:
            try:
                dblnk.lockdb()
                have_lock = True
            except PermissionDenied:
                pass

        try:
            if not dblnk.exists():
                # unmerged by a concurrent process
                continue
            iuse, use, restrict = vardb.aux_get(cpv, ["IUSE", "USE", "RESTRICT"])
            iuse = [x.lstrip("+-") for x in iuse.split()]
            use = use.split()
            try:
                restrict = use_reduce(restrict, uselist=use, flat=True)
            except InvalidDependString as e:
                eout.eerror(
                    "Invalid RESTRICT metadata " + f"for '{cpv}': {str(e)}; skipping"
                )
                del e
                continue
            if "bindist" in iuse and "bindist" not in use:
                eout.ewarn(f"{cpv}: package was emerged with USE=-bindist!")
                eout.ewarn(f"{cpv}: it might not be legal to redistribute this.")
            elif "bindist" in restrict:
                eout.ewarn(f"{cpv}: package has RESTRICT=bindist!")
                eout.ewarn(f"{cpv}: it might not be legal to redistribute this.")
            eout.ebegin(f"Building package for {cpv}")
            pkgs_for_arg += 1
            existing_metadata = dict(
                zip(fix_metadata_keys, vardb.aux_get(cpv, fix_metadata_keys))
            )
            category, pf = portage.catsplit(cpv)
            required_metadata = {}
            required_metadata["CATEGORY"] = category
            required_metadata["PF"] = pf
            update_metadata = {}
            for k, v in required_metadata.items():
                if v != existing_metadata[k]:
                    update_metadata[k] = v
            if update_metadata:
                vardb.aux_update(cpv, update_metadata)

            binpkg_format = settings.get(
                "BINPKG_FORMAT", SUPPORTED_GENTOO_BINPKG_FORMATS[0]
            )
            if binpkg_format == "xpak":
                xpdata = xpak.xpak(dblnk.dbdir)
                binpkg_tmpfile = os.path.join(
                    bintree.pkgdir, cpv + ".tbz2." + str(portage.getpid())
                )
                ensure_dirs(os.path.dirname(binpkg_tmpfile))
                binpkg_compression = settings.get("BINPKG_COMPRESS", "bzip2")
                try:
                    compression = _compressors[binpkg_compression]
                except KeyError as e:
                    if binpkg_compression:
                        eout.eerror(
                            f"Invalid or unsupported compression method: {e.args[0]}"
                        )
                        return 1
                    # Empty BINPKG_COMPRESS disables compression.
                    binpkg_compression = "none"
                    compression = {
                        "compress": "cat",
                        "package": "sys-apps/coreutils",
                    }

                if (
                    settings.get(
                        f"BINPKG_COMPRESS_FLAGS_{binpkg_compression.upper()}", None
                    )
                    is not None
                ):
                    compression["compress"] = compression["compress"].replace(
                        "${BINPKG_COMPRESS_FLAGS}",
                        f"${{BINPKG_COMPRESS_FLAGS_{binpkg_compression.upper()}}}",
                    )

                try:
                    compression_binary = shlex_split(
                        varexpand(compression["compress"], mydict=settings)
                    )[0]
                except IndexError as e:
                    eout.eerror(
                        f"Invalid or unsupported compression method: {e.args[0]}"
                    )
                    return 1
                if find_binary(compression_binary) is None:
                    missing_package = compression["package"]
                    eout.eerror(
                        f"File compression unsupported {binpkg_compression} (missing package: {missing_package})"
                    )
                    return 1

                cmd = compression["compress"].replace(
                    "{JOBS}",
                    str(makeopts_to_job_count(settings.get("MAKEOPTS", "1"))),
                )
                cmd = shlex_split(varexpand(cmd, mydict=settings))
                # Filter empty elements that make Popen fail
                cmd = [x for x in cmd if x != ""]
                with open(binpkg_tmpfile, "wb") as fobj:
                    proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=fobj)
                    excluded_config_files = dblnk.quickpkg(
                        proc.stdin,
                        include_config=include_config,
                        include_unmodified_config=include_unmodified_config,
                    )
                    proc.stdin.close()
                    if proc.wait() != os.EX_OK:
                        eout.eend(1)
                        eout.eerror(f"Compressor failed for package {cpv}")
                        retval |= 1
                        try:
                            os.unlink(binpkg_tmpfile)
                        except OSError as e:
                            if e.errno not in (errno.ENOENT, errno.ESTALE):
                                raise
                        continue
                xpak.tbz2(binpkg_tmpfile).recompose_mem(xpdata)
            elif binpkg_format == "gpkg":
                metadata = gpkg.gpkg(settings)._generate_metadata_from_dir(dblnk.dbdir)
                binpkg_tmpfile = os.path.join(
                    bintree.pkgdir, cpv + ".gpkg.tar." + str(os.getpid())
                )
                ensure_dirs(os.path.dirname(binpkg_tmpfile))
                excluded_config_files = dblnk.quickpkg(
                    binpkg_tmpfile,
                    metadata,
                    include_config=include_config,
                    include_unmodified_config=include_unmodified_config,
                )
            else:
                raise InvalidBinaryPackageFormat(binpkg_format)
        except CompressorOperationFailed as e:
            eout.eerror(f"Compressor operation failed")
            os.remove(binpkg_tmpfile)
            binpkg_tmpfile = None
        finally:
            if have_lock:
                dblnk.unlockdb()
        if binpkg_tmpfile is not None:
            pkg_info = bintree.inject(cpv, current_pkg_path=binpkg_tmpfile)
            # The pkg_info value ensures that the following getname call
            # returns the correct path when FEATURES=binpkg-multi-instance
            # is enabled, but fallback to cpv in case the inject call
            # returned None due to some kind of failure.
            binpkg_path = bintree.getname(pkg_info or cpv)
            try:
                s = os.stat(binpkg_path)
            except OSError:
                s = None

        if binpkg_tmpfile is None or s is None or pkg_info is None:
            # Sanity check, shouldn't happen normally.
            eout.eend(1)
            eout.eerror(f"Failed to create package: '{cpv}'")
            retval |= 1
        else:
            eout.eend(0)
            infos["successes"].append((cpv, s.st_size))
            infos["config_files_excluded"] += len(excluded_config_files)
            for filename in excluded_config_files:
                eout.ewarn(f"Excluded config: '{filename}'")
    if not pkgs_for_arg:
        eout.eerror("Could not find anything " + f"to match '{arg}'; skipping")
        infos["missing"].append(arg)
        retval |= 1
    return retval


def quickpkg_set(options, infos, arg, eout):
    eroot = portage.settings["EROOT"]
    trees = portage.db[eroot]
    vartree = trees["vartree"]

    settings = vartree.settings
    settings._init_dirs()
    setconfig = load_default_config(settings, trees)
    sets = setconfig.getSets()

    set_name = arg[1:]
    if not set_name in sets:
        eout.eerror(f"Package set not found: '{arg}'; skipping")
        infos["missing"].append(arg)
        return 1

    try:
        atoms = setconfig.getSetAtoms(set_name)
    except PackageSetNotFound as e:
        eout.eerror(
            f"Failed to process package set '{set_name}' because "
            + f"it contains the non-existent package set '{e}'; skipping"
        )
        infos["missing"].append(arg)
        return 1
    retval = os.EX_OK
    for atom in atoms:
        retval |= quickpkg_atom(options, infos, atom, eout)
    return retval


def quickpkg_extended_atom(options, infos, atom, eout):
    eroot = portage.settings["EROOT"]
    trees = portage.db[eroot]
    vartree = trees["vartree"]
    vardb = vartree.dbapi

    require_metadata = atom.slot or atom.repo
    atoms = []
    for cpv in vardb.cpv_all():
        cpv_atom = Atom(f"={cpv}")

        if atom == "*/*":
            atoms.append(cpv_atom)
            continue

        if not portage.match_from_list(atom, [cpv]):
            continue

        if require_metadata:
            try:
                cpv = vardb._pkg_str(cpv, atom.repo)
            except (KeyError, InvalidData):
                continue
            if not portage.match_from_list(atom, [cpv]):
                continue

        atoms.append(cpv_atom)

    for atom in atoms:
        quickpkg_atom(options, infos, atom, eout)


def quickpkg_main(options, args, eout):
    eroot = portage.settings["EROOT"]
    trees = portage.db[eroot]
    bintree = trees["bintree"]

    try:
        ensure_dirs(bintree.pkgdir)
    except portage.exception.PortageException:
        pass
    if not os.access(bintree.pkgdir, os.W_OK):
        eout.eerror(f"No write access to '{bintree.pkgdir}'")
        return errno.EACCES

    if "xattr" in portage.settings.features and not _xattr.XATTRS_WORKS:
        eout.eerror(
            "No xattr support library was found, " "so xattrs will not be preserved!"
        )
        portage.settings.unlock()
        portage.settings.features.remove("xattr")
        portage.settings.lock()

    infos = {}
    infos["successes"] = []
    infos["missing"] = []
    infos["config_files_excluded"] = 0
    for arg in args:
        if arg[0] == SETPREFIX:
            quickpkg_set(options, infos, arg, eout)
            continue
        try:
            atom = Atom(arg, allow_wildcard=True, allow_repo=True)
        except (InvalidAtom, InvalidData):
            # maybe it's valid but missing category (requires dep_expand)
            quickpkg_atom(options, infos, arg, eout)
        else:
            if atom.extended_syntax:
                quickpkg_extended_atom(options, infos, atom, eout)
            else:
                quickpkg_atom(options, infos, atom, eout)

    if not infos["successes"]:
        eout.eerror("No packages found")
        return 1
    print()
    eout.einfo(f"Packages now in '{bintree.pkgdir}':")
    units = {10: "K", 20: "M", 30: "G", 40: "T", 50: "P", 60: "E", 70: "Z", 80: "Y"}
    for cpv, size in infos["successes"]:
        if not size:
            # avoid OverflowError in math.log()
            size_str = "0"
        else:
            power_of_2 = math.log(size, 2)
            power_of_2 = 10 * (power_of_2 // 10)
            unit = units.get(power_of_2)
            if unit:
                size = float(size) / (2**power_of_2)
                size_str = f"{size:.1f}"
                if len(size_str) > 4:
                    # emulate `du -h`, don't show too many sig figs
                    size_str = str(int(size))
                size_str += unit
            else:
                size_str = str(size)
        eout.einfo(f"{cpv}: {size_str}")
    if infos["config_files_excluded"]:
        print()
        eout.ewarn(f"Excluded config files: {infos['config_files_excluded']}")
        eout.ewarn("See --help if you would like to include config files.")
    if infos["missing"]:
        print()
        eout.ewarn("The following packages could not be found:")
        eout.ewarn(" ".join(infos["missing"]))
        return 2
    return os.EX_OK


if __name__ == "__main__":
    usage = "quickpkg [options] <list of package atoms or package sets>"
    parser = argparse.ArgumentParser(usage=usage)
    parser.add_argument(
        "--umask",
        default="0077",
        help="umask used during package creation (default is 0077)",
    )
    parser.add_argument(
        "--ignore-default-opts",
        action="store_true",
        help="do not use the QUICKPKG_DEFAULT_OPTS environment variable",
    )
    parser.add_argument(
        "--include-config",
        choices=["y", "n"],
        default="n",
        metavar="<y|n>",
        help="include all files protected by CONFIG_PROTECT (as a security precaution, default is 'n')",
    )
    parser.add_argument(
        "--include-unmodified-config",
        choices=["y", "n"],
        default="n",
        metavar="<y|n>",
        help="include files protected by CONFIG_PROTECT that have not been modified since installation (as a security precaution, default is 'n')",
    )
    options, args = parser.parse_known_args(sys.argv[1:])
    if not options.ignore_default_opts:
        default_opts = shlex_split(portage.settings.get("QUICKPKG_DEFAULT_OPTS", ""))
        options, args = parser.parse_known_args(default_opts + sys.argv[1:])
    if not args:
        parser.error("no packages atoms given")
    try:
        umask = int(options.umask, 8)
    except ValueError:
        parser.error(f"invalid umask: {options.umask}")
    # We need to ensure a sane umask for the packages that will be created.
    old_umask = os.umask(umask)
    eout = portage.output.EOutput()

    def sigwinch_handler(signum, frame):
        lines, eout.term_columns = portage.output.get_term_size()

    gpg = None
    if portage.settings.get("BINPKG_GPG_SIGNING_KEY", None):
        gpg = GPG(portage.settings)
        gpg.unlock()

    signal.signal(signal.SIGWINCH, sigwinch_handler)
    try:
        retval = quickpkg_main(options, args, eout)
    finally:
        os.umask(old_umask)
        signal.signal(signal.SIGWINCH, signal.SIG_DFL)
        if gpg is not None:
            gpg.stop()
        global_event_loop().close()
    sys.exit(retval)
